#include"algorithm.hpp"

algorithm::algorithm(double delta,double T0,int size,output& out0):out(out0){
  dt=delta;
  i=0;
  T=T0;
  //rng.seed(std::time(0));
  rng.set();
  ensemble.resize(size);
  ensemble2.resize(size);
  for(int ii=0;ii<size;ii++){
    ensemble(ii)=0;
    ensemble2(ii)=0;
  }
}

// algorithm::algorithm(double delta,double T0,int size,output& out0,
// 		     unsigned seed):out(out0){
//   dt=delta;
//   i=0;
//   T=T0;
//   rng.seed(seed);
//   ensemble.resize(size);
//   ensemble2.resize(size);
//   for(int ii=0;ii<size;ii++){
//     ensemble(ii)=0;
//     ensemble2(ii)=0;
//   }
// }

double algorithm::get_time(){
    return i*dt;
}

void algorithm::update_time(){
  i+=1;
}

vector<double> algorithm::probability(int n){
  //n=number of states in effective_ensemble
  vector<double> r(n);
  for(int ii=0;ii<n;ii++){
    r(ii)=0;
    for(unsigned jj=0;jj<ensemble.size();jj++){
      if(ensemble[jj]==ii){
	r(ii)+=1;
      }
    }
  }
  
  if(sum(r)!=ensemble.size()){
    std::cout<<"Error,ensemble size mismatch!!!"<<std::endl;
  }
  return r/double(ensemble.size());
}

vector<double> algorithm::sigma(effective_ensemble& eff){
  int num_states=eff.get_number_of_states();
  int dim=eff.get_dim();
  int N=ensemble.size();
  identity_matrix<double> I(dim);
  matrix<std::complex<double> > rho=eff.rho();
  vector<double> s(dim*dim);
  for(int ii=0;ii<dim*dim;ii++){
    s(ii)=0;
  }
  for(int ii=0;ii<dim;ii++){
    for(int jj=0;jj<dim;jj++){
      for(int kk=0;kk<num_states;kk++){
	s(ii*dim+jj)=+std::pow(std::abs((std::pow(norm_2(prod(outer_prod(column(I,ii),column(I,jj)),eff.get_state(kk))),2)*eff.P(kk))-
					rho(ii,jj)/double(N)),2);
	  
	  
      }
    }
  }
  for(int ii=0;ii<dim*dim;ii++){
    s(ii)=std::sqrt((double(N)/double(N-1))*s(ii));
  }
  return s;
}


void algorithm::make_transition(int source,vector<int> target,
				vector<double> p){
  double ps=sum(p);
  for(unsigned ii=0;ii<ensemble.size();ii++){
    if(ensemble(ii)==source){
      if(rng.sample()<=ps){
	//hyppy
	
	if(target.size()==1){
	  ensemble2(ii)=target(0);
	}
	
	else{
	  ensemble2(ii)=target(rng.sample_int_from(p));
	}
      }
    }
  }
}
// double algorithm::sample(){
//   static boost::uniform_01<boost::mt19937> dist(rng);
//   return dist();
// }


void algorithm::det_evol(effective_ensemble& eff,channels& chans){
  int n_c=chans.get_number_of_channels();
  int n_s=eff.get_number_of_states();
  double time=get_time();
  
  //Effective Hamiltonian
  matrix<std::complex<double> > H;
  for(int ii=0;ii<n_c;ii++){
    if(ii==0){
      H=chans.get_rate(ii,time)*prod(herm(chans.get_op(ii)),chans.get_op(ii));
    }
    else
      H=H+chans.get_rate(ii,time)*prod(herm(chans.get_op(ii)),chans.get_op(ii));
  }
  H*=std::complex<double>(0.0,-0.5);
  
  //Evolution of the states
  vector<std::complex<double> > vec;
  for(int ii=0;ii<n_s;ii++){
    vec=eff.get_state(ii);
    vec=vec+std::complex<double>(0.0,-dt)*prod(H,vec);
    vec=vec/norm_2(vec);
    
    eff.set_state(ii,vec);
  }
  
};

void algorithm::jumps(effective_ensemble& eff,channels& chans, 
		      transition_map tmap[]){
  
  int n_s=eff.get_number_of_states();
  int n_c=chans.get_number_of_channels();
  double time=get_time();
  vector<double> p_jump(n_c*n_s);
  vector<int> target(n_c*n_s);
  left_const_iterator il,ile;
  right_const_iterator ir,ire;
  
  for(int hh=0;hh<n_s;hh++){//iterate over states
    int num=0;
    if(hh!=0){
      p_jump.resize(n_c*n_s);
      target.resize(n_c*n_s);
    }
    for(int ii=0;ii<n_c;ii++){//iterate over decay channels
      
      if(chans.get_rate(ii,time)>=0){//positive channel
	//find if jump away from state hh via channel ii is possible
	
	lpair targs=tmap[ii].left.equal_range(hh);
	if(targs.first!=tmap[ii].left.end()){
	  for(left_const_iterator il=targs.first;il!=targs.second;il++){
	    
	    p_jump(num)=dt*chans.get_rate(ii,time)*
	      pow(norm_2(prod(chans.get_op(ii),eff.get_state(hh))),2);
	    
	    target(num)=il->second;
	    
	    num+=1;
	  }
	}
      }
      else if(chans.get_rate(ii,time)<0){//negative channel
	//find if jump away from state hh via channel ii is possible
	rpair targs=tmap[ii].right.equal_range(hh);
	if(targs.first!=tmap[ii].right.end()){
	  for(right_const_iterator ir=targs.first;ir!=targs.second;ir++){
	    
	    p_jump(num)=dt*std::abs(chans.get_rate(ii,time))*eff.P(ir->second)/
	      eff.P(hh)*
	      pow(norm_2(prod(chans.get_op(ii),eff.get_state(ir->second))),2);
	    
	    target(num)=ir->second;
	    
	    num+=1;
	  }
	}
      }
    } 
    //if num == 0 then there is no transitions from state hh
    if(num!=0){
      p_jump.resize(num,true);
      target.resize(num,true);
      make_transition(hh,target,p_jump);
    }
    
    ensemble=ensemble2;
    eff.set_P(probability(n_s));
  }
}
void algorithm::apply(effective_ensemble& eff,channels& chans,
		      transition_map tmap[]){
  // output out(eff.get_number_of_states(),chans.get_number_of_channels(),
  // 	     ensemble.size(),ceil(T/dt));
  
  out.set_ensemble(ensemble,0);
  out.set_P(eff.p(),0);
  out.set_rho_vec(eff.rho_vec(),0);
  out.set_sigma(sigma(eff),0);
  for(int it=1;it<ceil(T/dt);it++){
    //make jumps
    jumps(eff,chans,tmap);
    //evolve deterministically and update time
    det_evol(eff,chans);
    //write to output variables
    out.set_ensemble(ensemble,it);
    out.set_P(eff.p(),it);
    out.set_rho_vec(eff.rho_vec(),it);
    out.set_sigma(sigma(eff),it);
    //Update time
    update_time();
  }
}

output algorithm::get_output(){
  return out;
}

void algorithm::print_to_files(){
  out.print_to_files();
}
